import rosbag
from sensor_msgs.msg import PointCloud2
from sensor_msgs import point_cloud2

# 输入和输出的ROS bag文件路径
input_bag_path = "/media/atakanoinu/MSDAD/[0]calibration/calibration.bag"
output_bag_path = "/media/atakanoinu/MSDAD/[0]calibration/calibration_overlay.bag"

# 要叠加的五帧点云的topic名称
input_topic = "/solid_lidar/point"

# 新的叠加后的点云的topic名称
output_topic = "/solid_lidar/point"

# 初始化输出ROS bag
output_bag = rosbag.Bag(output_bag_path, 'w')
combined_points = []

with rosbag.Bag(input_bag_path, 'r') as input_bag:
    # 记录计数器
    count = 0

    # 遍历输入bag中的消息
    for topic, msg, t in input_bag.read_messages():
        if topic == input_topic:
            # 将点云数据添加到叠加的点云列表中
            for point in point_cloud2.read_points(msg, field_names=("x", "y", "z"), skip_nans=True):
                combined_points.append(point)

            # 增加计数器
            count += 1

            # 当收集到五帧点云时进行叠加
            if count == 5:
                # 创建新的点云消息
                combined_pc = point_cloud2.create_cloud_xyz32(header=msg.header, points=combined_points)

                # 将叠加后的点云写入输出ROS bag
                output_bag.write(output_topic, combined_pc, t)

                # 重置计数器和点云列表
                count = 0
                combined_points = []
        else:
        # 其他topic不变
            output_bag.write(topic, msg, t)


# 关闭输出ROS bag
output_bag.close()